package com.wk.filter;

import com.google.common.util.concurrent.RateLimiter;
import com.netflix.zuul.ZuulFilter;
import com.netflix.zuul.context.RequestContext;
import com.netflix.zuul.exception.ZuulException;
import org.springframework.cloud.netflix.zuul.filters.support.FilterConstants;
import org.springframework.http.HttpStatus;
import org.springframework.stereotype.Component;

/**
 * 使用guava令牌桶在网关限流
 */
@Component
public class LimitFilter extends ZuulFilter {


	@Override
	public String filterType() {
		//限流放在最前面执行
		return FilterConstants.PRE_TYPE;
	}

	@Override
	public int filterOrder() {
		return 0;
	}

	@Override
	public boolean shouldFilter() {
		return true;
	}

	//设置令牌生成器生成速率，每秒生成2个，如果是小数，如0.1，则10秒生成1个
	private static final RateLimiter LIMITER = RateLimiter.create(2);

	@Override
	public Object run() throws ZuulException {

		RequestContext context = RequestContext.getCurrentContext();

		if (LIMITER.tryAcquire()) {
			return null;		//如果获取到令牌继续执行
		}else {
			//获取不到令牌则不向后执行，并返回一个状态码
			//setSendZuulResponse设置为false则不向pre后的过滤器，这里要求后面所有的过滤器都不执行，
			// 设置一个自定义属性，在后面过滤器的shouldFilter方法中添加判断，如果limit是false则不执行
			context.set("limit",false);
			context.setSendZuulResponse(false);
			context.setResponseStatusCode(HttpStatus.TOO_MANY_REQUESTS.value());
		}

		return null;
	}
}
